import cv2 as cv
import numpy as np
from cv2.mat_wrapper import Mat


def getCenter(x, y, w, h):
    x1 = w / 2
    y1 = h / 2
    cx = x + x1
    cy = y + y1
    return cx, cy


cap = cv.VideoCapture("418.mp4")
backsub = cv.createBackgroundSubtractorMOG2()
#车辆形态大小
car_w = 30
car_h = 30
#统计分界线位置
line_h = 700
#车数量
carno = 0
cars = []
#线周边偏移量
line_offset = 10

# 形态学
kernel = cv.getStructuringElement(cv.MORPH_RECT, (5, 5))
while True:
    ret, frame = cap.read()
    # print(frame.shape)
    if ret:
        gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
        blur = cv.GaussianBlur(gray, (3, 3), 20)
        # 去背景
        fgmask = backsub.apply(blur, None, 0.01)
        # 腐蚀
        erode = cv.erode(fgmask, kernel)
        # 膨胀
        dilate = cv.dilate(erode, kernel, iterations=2)
        # 闭操作，去掉物体内的小块
        close = cv.morphologyEx(dilate, cv.MORPH_CLOSE, kernel)
        close = cv.morphologyEx(close, cv.MORPH_CLOSE, kernel)
        # 查找轮廓
        cnts, h = cv.findContours(close, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)
        for (i, c) in enumerate(cnts):
            (x, y, w, h) = cv.boundingRect(c)
            isValid = (w >= car_w) and (h >= car_h)
            if not isValid:
                continue

            cv.line(frame, (10, line_h), (1480, line_h), (0, 0, 255), 3)
            cv.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0),2)

            cpoint = getCenter(x, y, w, h)
            cars.append(cpoint)
            for (x1, y1) in (cars):
                print(cars)
                if (y1 < line_h - line_offset) and (y1 > line_h + line_offset):
                    carno += 1
                    cars.remove((x,y))
                    print(carno)



        # cv.imshow("video", fgmask)
        # cv.imshow("close", close)
        cv.imshow("frame", frame)
        key = cv.waitKey(50)
        if key == ord("q"):
            print("退出。。。")
            break
    else:
        break

cap.release()
cv.destroyAllWindows()
